#define JPEG_INTERNALS
#include "jinclude.h"
#include "jpeglib.h"

typedef struct {
  struct jpeg_color_converter pub; /* public fields */
  JPEG_INT32 * rgb_ycc_tab;		/* => table for RGB to YCbCr conversion */
} my_color_converter;

typedef my_color_converter * my_cconvert_ptr;

#define SCALEBITS	16	/* speediest right-shift on some machines */
#define CBCR_OFFSET	((JPEG_INT32) CENTERJSAMPLE << SCALEBITS)
#define ONE_HALF	((JPEG_INT32) 1 << (SCALEBITS-1))
#define FIX(x)		((JPEG_INT32) ((x) * (1L<<SCALEBITS) + 0.5))
#define R_Y_OFF		0			/* offset to R => Y section */
#define G_Y_OFF		(1*(MAXJSAMPLE+1))	/* offset to G => Y section */
#define B_Y_OFF		(2*(MAXJSAMPLE+1))	/* etc. */
#define R_CB_OFF	(3*(MAXJSAMPLE+1))
#define G_CB_OFF	(4*(MAXJSAMPLE+1))
#define B_CB_OFF	(5*(MAXJSAMPLE+1))
#define R_CR_OFF	B_CB_OFF		/* B=>Cb, R=>Cr are the same */
#define G_CR_OFF	(6*(MAXJSAMPLE+1))
#define B_CR_OFF	(7*(MAXJSAMPLE+1))
#define TABLE_SIZE	(8*(MAXJSAMPLE+1))

METHODDEF( void )
rgb_ycc_start( j_compress_ptr cinfo ) {
  my_cconvert_ptr cconvert = ( my_cconvert_ptr ) cinfo->cconvert;
  JPEG_INT32 * rgb_ycc_tab;
  JPEG_INT32 i;
  /* Allocate and fill in the conversion tables. */
  cconvert->rgb_ycc_tab = rgb_ycc_tab = ( JPEG_INT32 * )
                                        ( *cinfo->mem->alloc_small )( ( j_common_ptr ) cinfo, JPOOL_IMAGE,
                                            ( TABLE_SIZE * SIZEOF( JPEG_INT32 ) ) );
  for( i = 0; i <= MAXJSAMPLE; i++ ) {
    rgb_ycc_tab[i + R_Y_OFF] = FIX( 0.29900 ) * i;
    rgb_ycc_tab[i + G_Y_OFF] = FIX( 0.58700 ) * i;
    rgb_ycc_tab[i + B_Y_OFF] = FIX( 0.11400 ) * i     + ONE_HALF;
    rgb_ycc_tab[i + R_CB_OFF] = ( -FIX( 0.16874 ) ) * i;
    rgb_ycc_tab[i + G_CB_OFF] = ( -FIX( 0.33126 ) ) * i;
    rgb_ycc_tab[i + B_CB_OFF] = FIX( 0.50000 ) * i    + CBCR_OFFSET + ONE_HALF - 1;
    rgb_ycc_tab[i + G_CR_OFF] = ( -FIX( 0.41869 ) ) * i;
    rgb_ycc_tab[i + B_CR_OFF] = ( -FIX( 0.08131 ) ) * i;
  }
}

METHODDEF( void ) rgb_ycc_convert( j_compress_ptr cinfo,
                                   JSAMPARRAY input_buf, JSAMPIMAGE output_buf,
                                   JDIMENSION output_row, int num_rows ) {
  my_cconvert_ptr cconvert = ( my_cconvert_ptr ) cinfo->cconvert;
  register int r, g, b;
  register JPEG_INT32 * ctab = cconvert->rgb_ycc_tab;
  register JSAMPROW inptr;
  register JSAMPROW outptr0, outptr1, outptr2;
  register JDIMENSION col;
  JDIMENSION num_cols = cinfo->image_width;
  while( --num_rows >= 0 ) {
    inptr = *input_buf++;
    outptr0 = output_buf[0][output_row];
    outptr1 = output_buf[1][output_row];
    outptr2 = output_buf[2][output_row];
    output_row++;
    for( col = 0; col < num_cols; col++ ) {
      r = GETJSAMPLE( inptr[RGB_RED] );
      g = GETJSAMPLE( inptr[RGB_GREEN] );
      b = GETJSAMPLE( inptr[RGB_BLUE] );
      inptr += RGB_PIXELSIZE;
      /* Y */
      outptr0[col] = ( JSAMPLE )
                     ( ( ctab[r + R_Y_OFF] + ctab[g + G_Y_OFF] + ctab[b + B_Y_OFF] )
                       >> SCALEBITS );
      /* Cb */
      outptr1[col] = ( JSAMPLE )
                     ( ( ctab[r + R_CB_OFF] + ctab[g + G_CB_OFF] + ctab[b + B_CB_OFF] )
                       >> SCALEBITS );
      /* Cr */
      outptr2[col] = ( JSAMPLE )
                     ( ( ctab[r + R_CR_OFF] + ctab[g + G_CR_OFF] + ctab[b + B_CR_OFF] )
                       >> SCALEBITS );
    }
  }
}

METHODDEF( void ) rgb_gray_convert( j_compress_ptr cinfo,
                                    JSAMPARRAY input_buf, JSAMPIMAGE output_buf,
                                    JDIMENSION output_row, int num_rows ) {
  my_cconvert_ptr cconvert = ( my_cconvert_ptr ) cinfo->cconvert;
  register int r, g, b;
  register JPEG_INT32 * ctab = cconvert->rgb_ycc_tab;
  register JSAMPROW inptr;
  register JSAMPROW outptr;
  register JDIMENSION col;
  JDIMENSION num_cols = cinfo->image_width;
  while( --num_rows >= 0 ) {
    inptr = *input_buf++;
    outptr = output_buf[0][output_row];
    output_row++;
    for( col = 0; col < num_cols; col++ ) {
      r = GETJSAMPLE( inptr[RGB_RED] );
      g = GETJSAMPLE( inptr[RGB_GREEN] );
      b = GETJSAMPLE( inptr[RGB_BLUE] );
      inptr += RGB_PIXELSIZE;
      /* Y */
      outptr[col] = ( JSAMPLE )
                    ( ( ctab[r + R_Y_OFF] + ctab[g + G_Y_OFF] + ctab[b + B_Y_OFF] )
                      >> SCALEBITS );
    }
  }
}

METHODDEF( void ) cmyk_ycck_convert( j_compress_ptr cinfo,
                                     JSAMPARRAY input_buf, JSAMPIMAGE output_buf,
                                     JDIMENSION output_row, int num_rows ) {
  my_cconvert_ptr cconvert = ( my_cconvert_ptr ) cinfo->cconvert;
  register int r, g, b;
  register JPEG_INT32 * ctab = cconvert->rgb_ycc_tab;
  register JSAMPROW inptr;
  register JSAMPROW outptr0, outptr1, outptr2, outptr3;
  register JDIMENSION col;
  JDIMENSION num_cols = cinfo->image_width;
  while( --num_rows >= 0 ) {
    inptr = *input_buf++;
    outptr0 = output_buf[0][output_row];
    outptr1 = output_buf[1][output_row];
    outptr2 = output_buf[2][output_row];
    outptr3 = output_buf[3][output_row];
    output_row++;
    for( col = 0; col < num_cols; col++ ) {
      r = MAXJSAMPLE - GETJSAMPLE( inptr[0] );
      g = MAXJSAMPLE - GETJSAMPLE( inptr[1] );
      b = MAXJSAMPLE - GETJSAMPLE( inptr[2] );
      /* K passes through as-is */
      outptr3[col] = inptr[3];	/* don't need GETJSAMPLE here */
      inptr += 4;
      /* If the inputs are 0..MAXJSAMPLE, the outputs of these equations
         must be too; we do not need an explicit range-limiting operation.
         Hence the value being shifted is never negative, and we don't
         need the general RIGHT_SHIFT macro.
      */
      /* Y */
      outptr0[col] = ( JSAMPLE )
                     ( ( ctab[r + R_Y_OFF] + ctab[g + G_Y_OFF] + ctab[b + B_Y_OFF] )
                       >> SCALEBITS );
      /* Cb */
      outptr1[col] = ( JSAMPLE )
                     ( ( ctab[r + R_CB_OFF] + ctab[g + G_CB_OFF] + ctab[b + B_CB_OFF] )
                       >> SCALEBITS );
      /* Cr */
      outptr2[col] = ( JSAMPLE )
                     ( ( ctab[r + R_CR_OFF] + ctab[g + G_CR_OFF] + ctab[b + B_CR_OFF] )
                       >> SCALEBITS );
    }
  }
}

METHODDEF( void )
grayscale_convert( j_compress_ptr cinfo,
                   JSAMPARRAY input_buf, JSAMPIMAGE output_buf,
                   JDIMENSION output_row, int num_rows ) {
  register JSAMPROW inptr;
  register JSAMPROW outptr;
  register JDIMENSION col;
  JDIMENSION num_cols = cinfo->image_width;
  int instride = cinfo->input_components;
  while( --num_rows >= 0 ) {
    inptr = *input_buf++;
    outptr = output_buf[0][output_row];
    output_row++;
    for( col = 0; col < num_cols; col++ ) {
      outptr[col] = inptr[0];	/* don't need GETJSAMPLE() here */
      inptr += instride;
    }
  }
}

METHODDEF( void )
null_convert( j_compress_ptr cinfo,
              JSAMPARRAY input_buf, JSAMPIMAGE output_buf,
              JDIMENSION output_row, int num_rows ) {
  register JSAMPROW inptr;
  register JSAMPROW outptr;
  register JDIMENSION col;
  register int ci;
  int nc = cinfo->num_components;
  JDIMENSION num_cols = cinfo->image_width;
  while( --num_rows >= 0 ) {
    /* It seems fastest to make a separate pass for each component. */
    for( ci = 0; ci < nc; ci++ ) {
      inptr = *input_buf;
      outptr = output_buf[ci][output_row];
      for( col = 0; col < num_cols; col++ ) {
        outptr[col] = inptr[ci]; /* don't need GETJSAMPLE() here */
        inptr += nc;
      }
    }
    input_buf++;
    output_row++;
  }
}

METHODDEF( void )
null_method( j_compress_ptr cinfo ) {
  /* no work needed */
}

void jinit_color_converter( j_compress_ptr cinfo ) {
  my_cconvert_ptr cconvert;
  cconvert = ( my_cconvert_ptr )
             ( *cinfo->mem->alloc_small )( ( j_common_ptr ) cinfo, JPOOL_IMAGE,
                 SIZEOF( my_color_converter ) );
  cinfo->cconvert = ( struct jpeg_color_converter * ) cconvert;
  /* set start_pass to null method until we find out differently */
  cconvert->pub.start_pass = null_method;
  /* Make sure input_components agrees with in_color_space */
  switch( cinfo->in_color_space ) {
    case JCS_GRAYSCALE:
      if( cinfo->input_components != 1 )
      { ERREXIT( cinfo, JERR_BAD_IN_COLORSPACE ); }
      break;
    case JCS_RGB:
      #if RGB_PIXELSIZE != 3
      if( cinfo->input_components != RGB_PIXELSIZE )
      { ERREXIT( cinfo, JERR_BAD_IN_COLORSPACE ); }
      break;
      #endif /* else share code with YCbCr */
    case JCS_YCbCr:
      if( cinfo->input_components != 3 )
      { ERREXIT( cinfo, JERR_BAD_IN_COLORSPACE ); }
      break;
    case JCS_CMYK:
    case JCS_YCCK:
      if( cinfo->input_components != 4 )
      { ERREXIT( cinfo, JERR_BAD_IN_COLORSPACE ); }
      break;
    default:			/* JCS_UNKNOWN can be anything */
      if( cinfo->input_components < 1 )
      { ERREXIT( cinfo, JERR_BAD_IN_COLORSPACE ); }
      break;
  }
  /* Check num_components, set conversion method based on requested space */
  switch( cinfo->jpeg_color_space ) {
    case JCS_GRAYSCALE:
      if( cinfo->num_components != 1 )
      { ERREXIT( cinfo, JERR_BAD_J_COLORSPACE ); }
      if( cinfo->in_color_space == JCS_GRAYSCALE )
      { cconvert->pub.color_convert = grayscale_convert; }
      else if( cinfo->in_color_space == JCS_RGB ) {
        cconvert->pub.start_pass = rgb_ycc_start;
        cconvert->pub.color_convert = rgb_gray_convert;
      } else if( cinfo->in_color_space == JCS_YCbCr )
      { cconvert->pub.color_convert = grayscale_convert; }
      else
      { ERREXIT( cinfo, JERR_CONVERSION_NOTIMPL ); }
      break;
    case JCS_RGB:
      if( cinfo->num_components != 3 )
      { ERREXIT( cinfo, JERR_BAD_J_COLORSPACE ); }
      if( cinfo->in_color_space == JCS_RGB && RGB_PIXELSIZE == 3 )
      { cconvert->pub.color_convert = null_convert; }
      else
      { ERREXIT( cinfo, JERR_CONVERSION_NOTIMPL ); }
      break;
    case JCS_YCbCr:
      if( cinfo->num_components != 3 )
      { ERREXIT( cinfo, JERR_BAD_J_COLORSPACE ); }
      if( cinfo->in_color_space == JCS_RGB ) {
        cconvert->pub.start_pass = rgb_ycc_start;
        cconvert->pub.color_convert = rgb_ycc_convert;
      } else if( cinfo->in_color_space == JCS_YCbCr )
      { cconvert->pub.color_convert = null_convert; }
      else
      { ERREXIT( cinfo, JERR_CONVERSION_NOTIMPL ); }
      break;
    case JCS_CMYK:
      if( cinfo->num_components != 4 )
      { ERREXIT( cinfo, JERR_BAD_J_COLORSPACE ); }
      if( cinfo->in_color_space == JCS_CMYK )
      { cconvert->pub.color_convert = null_convert; }
      else
      { ERREXIT( cinfo, JERR_CONVERSION_NOTIMPL ); }
      break;
    case JCS_YCCK:
      if( cinfo->num_components != 4 )
      { ERREXIT( cinfo, JERR_BAD_J_COLORSPACE ); }
      if( cinfo->in_color_space == JCS_CMYK ) {
        cconvert->pub.start_pass = rgb_ycc_start;
        cconvert->pub.color_convert = cmyk_ycck_convert;
      } else if( cinfo->in_color_space == JCS_YCCK )
      { cconvert->pub.color_convert = null_convert; }
      else
      { ERREXIT( cinfo, JERR_CONVERSION_NOTIMPL ); }
      break;
    default:			/* allow null conversion of JCS_UNKNOWN */
      if( cinfo->jpeg_color_space != cinfo->in_color_space ||
          cinfo->num_components != cinfo->input_components )
      { ERREXIT( cinfo, JERR_CONVERSION_NOTIMPL ); }
      cconvert->pub.color_convert = null_convert;
      break;
  }
}
